{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "intro_md"
      },
      "source": [
        "# Dataset Utilities"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "desc_md"
      },
      "source": [
        "The `gtsam/slam/dataset.h` header provides utility functions for loading and saving factor graph data, particularly in formats commonly used in the SLAM community like TORO ('.graph') and g2o ('.g2o').\n",
        "\n",
        "Key functions include:\n",
        "*   `findExampleDataFile(name)`: Locates example dataset files distributed with GTSAM.\n",
        "*   `load2D(filename, ...)`: Loads a 2D pose graph (Poses and BetweenFactors) from TORO or g2o format.\n",
        "*   `load3D(filename)`: Loads a 3D pose graph (currently simpler than `load2D`, assumes specific format).\n",
        "*   `readG2o(filename, is3D)`: A more general function to read g2o files containing various factor and variable types (2D/3D poses, landmarks, measurements).\n",
        "*   `writeG2o(graph, initialEstimate, filename)`: Saves a factor graph and values to the g2o format.\n",
        "*   `parseVariables<T>`/`parseMeasurements<T>`/`parseFactors<T>`: Lower-level parsing functions (less commonly used directly)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "colab_badge_md"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/borglab/gtsam/blob/develop/gtsam/slam/doc/dataset.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pip_code",
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "%pip install --quiet gtsam-develop"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "imports_code"
      },
      "outputs": [],
      "source": [
        "import gtsam\n",
        "import gtsam.utils.plot as gtsam_plot\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import os\n",
        "\n",
        "# Ensure the example datasets are available\n",
        "# In Colab/pip install, they are usually included.\n",
        "# If running locally from build, findExampleDataFile works.\n",
        "# If running locally from install without examples, this might fail.\n",
        "try:\n",
        "    # Check for a common example file\n",
        "    gtsam.findExampleDataFile(\"w100.graph\")\n",
        "    HAVE_EXAMPLES = True\n",
        "except RuntimeError:\n",
        "    print(\"Warning: Example datasets not found.\")\n",
        "    print(\"Try running from build directory or installing examples.\")\n",
        "    HAVE_EXAMPLES = False\n",
        "\n",
        "# Create dummy files for writing examples if needed\n",
        "if not os.path.exists(\"output\"):\n",
        "    os.makedirs(\"output\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "find_header_md"
      },
      "source": [
        "## Finding Example Datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "find_example_code"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Found w100.graph at: c:\\Users\\porte\\miniconda3\\envs\\gtsam\\Lib\\site-packages\\gtsam\\Data\\w100.graph\n",
            "Found dlr.g2o at: None\n"
          ]
        }
      ],
      "source": [
        "if HAVE_EXAMPLES:\n",
        "    try:\n",
        "        w100_path = gtsam.findExampleDataFile(\"w100.graph\")\n",
        "        print(f\"Found w100.graph at: {w100_path}\")\n",
        "        dlr_path = gtsam.findExampleDataFile(\"dlr.g2o\")\n",
        "        print(f\"Found dlr.g2o at: {dlr_path}\")\n",
        "    except RuntimeError as e:\n",
        "        print(f\"Error finding example file: {e}\")\n",
        "else:\n",
        "    print(\"Skipping findExampleDataFile test as examples are not available.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "load2d_header_md"
      },
      "source": [
        "## Loading 2D Datasets (`load2D`)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "load2d_example_code"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Loaded c:\\Users\\porte\\miniconda3\\envs\\gtsam\\Lib\\site-packages\\gtsam\\Data\\w100.graph:\n",
            "  Graph size: 300\n",
            "  Initial estimate keys: 100\n"
          ]
        }
      ],
      "source": [
        "if HAVE_EXAMPLES:\n",
        "    # Load TORO 2D file\n",
        "    graph_2d, initial_2d = gtsam.load2D(w100_path)\n",
        "    print(f\"\\nLoaded {w100_path}:\")\n",
        "    print(f\"  Graph size: {graph_2d.size()}\")\n",
        "    print(f\"  Initial estimate keys: {len(initial_2d.keys())}\")\n",
        "\n",
        "    # Plot initial estimate (optional)\n",
        "    for key in initial_2d.keys():\n",
        "         gtsam_plot.plot_pose2(0, initial_2d.atPose2(key))\n",
        "    plt.title(\"Initial Estimate from w100.graph\")\n",
        "    plt.axis('equal')\n",
        "    # plt.show() # Uncomment to display plot\n",
        "    plt.close() # Close plot to prevent display in output\n",
        "else:\n",
        "    print(\"\\nSkipping load2D test.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g2o_header_md"
      },
      "source": [
        "## Loading/Saving G2O Files (`readG2o`, `writeG2o`)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g2o_desc_md"
      },
      "source": [
        "`readG2o` can handle both 2D and 3D datasets and various factor types defined in the g2o format.\n",
        "`writeG2o` saves a `NonlinearFactorGraph` and `Values` into a g2o file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "g2o_example_code"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Error processing None: readG2o(): incompatible function arguments. The following argument types are supported:\n",
            "    1. (filename: str, is3D: bool = False, kernelFunctionType: gtsam.gtsam.KernelFunctionType = <KernelFunctionType.KernelFunctionTypeNONE: 0>) -> tuple[gtsam.gtsam.NonlinearFactorGraph, gtsam.gtsam.Values]\n",
            "\n",
            "Invoked with: None; kwargs: is3D=True\n"
          ]
        }
      ],
      "source": [
        "if HAVE_EXAMPLES:\n",
        "   # Load a 3D g2o file\n",
        "   try:\n",
        "      graph_3d, initial_3d = gtsam.readG2o(dlr_path, is3D=True)\n",
        "      print(f\"\\nLoaded {dlr_path} (3D):\")\n",
        "      print(f\"  Graph size: {graph_3d.size()}\")\n",
        "      print(f\"  Initial estimate keys: {len(initial_3d.keys())}\")\n",
        "      # You could optimize graph_3d with initial_3d here...\n",
        "      print(\"Optimization skipped for brevity.\")\n",
        "      optimized_values = initial_3d # Use initial for demo write\n",
        "\n",
        "      # Save the graph and values back to a g2o file\n",
        "      output_g2o_file = os.path.join(\"output\", \"dlr_rewrite.g2o\")\n",
        "      gtsam.writeG2o(graph_3d, optimized_values, output_g2o_file)\n",
        "      print(f\"  Saved graph and values to {output_g2o_file}\")\n",
        "      # Clean up dummy file\n",
        "      # os.remove(output_g2o_file)\n",
        "      # os.rmdir(\"output\")\n",
        "   except (RuntimeError, TypeError) as e:\n",
        "      print(f\"Error processing {dlr_path}: {e}\")\n",
        "else:\n",
        "    print(\"\\nSkipping readG2o/writeG2o test.\")"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "gtsam",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.13.1"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
